import os

from langchain_chroma import Chroma
from langchain_community.document_loaders import PyPDFLoader, PyMuPDFLoader
from langchain_community.embeddings import DashScopeEmbeddings
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_openai import ChatOpenAI
from langchain_text_splitters import RecursiveCharacterTextSplitter

os.environ['LANGCHAIN_TRACING_V2'] = 'true'
os.environ['LANGCHAIN_PROJECT'] = 'LLMDEMO'
os.environ['LANGCHAIN_API_KEY'] = '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'
# 初始化带模板支持的LLM
# pip install redis
llm = ChatOpenAI(
    api_key="sk-a3f7718fb81f43b2915f0a6483b6661b",
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model="llama-4-scout-17b-16e-instruct",  # 此处以qwen-plus为例，您可按需更换模型名称。模型列表：https://help.aliyun.com/zh/model-studio/getting-started/models
    # other params...
)
#pip install pypdf
# pip install langchain-chroma
# pip install chroma-hnswlib==0.7.5 chromadb==0.5.4 版本不兼容报错问题
# pip install dashscope

# RAG流程 加载文档(PDF、CSV等) 向量化文档 相似性检索
def load_pdf(pdf_path):
    """使用PyMuPDF加载PDF文件"""
    loader = PyMuPDFLoader(pdf_path)
    documents = loader.load()
    return documents

# 2. 分割文本
def split_documents(documents):
    """使用LangChain的递归字符文本分割器"""
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=512,          # 每个块的大小
        chunk_overlap=50,        # 块之间的重叠
        length_function=len,     # 长度计算函数
        add_start_index=True,    # 添加起始索引
    )
    splits = text_splitter.split_documents(documents)
    print(f"已将文档分割为 {len(splits)} 个块")
    return splits


# 3. 创建向量存储
def create_vectorstore(splits, persist_directory="chroma_db"):
    """创建Chroma向量存储"""
    # 初始化嵌入模型
    model_name = "llama-4-scout-17b-16e-instruct"
    embeddings = DashScopeEmbeddings()

    # 创建并返回Chroma向量存储
    vectordb = Chroma.from_documents(
        documents=splits,
        embedding=embeddings,
        persist_directory=persist_directory
    )

    print(f"向量存储已保存到 {persist_directory}")
    return vectordb
# 1. 加载PDF
print("正在加载PDF文档...")
documents = load_pdf(r"D:\pythonWork\python_demo\pdf\问答题.pdf")
text_spliter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=100, length_function=len,add_start_index=True)
# 2. 分割文本
print("正在分割文本...")
splits = split_documents(documents)
# 3. 创建向量存储
print("正在创建向量存储...")
vectordb = create_vectorstore(splits)
#加载文档并把向量数据库保存到本地，保存在embedding文件夹下
retriever = RunnableLambda(vectordb.similarity_search).bind(k=1)
print(retriever)
prompt_template = """
你是一个智能问答机器人
你的任务是根据下述给定的已知信息回答用户问题,保证问答答案在50字以内。
确保你的回复完全依据一下已知信息，不要编造答案。
如果下述已知信息不足以回答用户的回复，请直接回复"我无法回答你的问题"
已知信息：
{info}
用户问题：
{input}
请用中文回答。
"""
prompt_template = ChatPromptTemplate.from_messages([("human",prompt_template)])

chain = {"input": RunnablePassthrough(), "info": retriever} | prompt_template | llm
print(chain.invoke("什么是JAVA？").content)
